﻿using System;
using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Data;
using System.Collections.Generic;

public class TerminalSession
{
    public int Id { get; set; }
    public string WinStationName { get; set; }
    public string UserName { get; set; }
    public string Domain { get; set; }
    public string Address { get; set; }
    public ConnectionState State { get; set; }
    public string UserDomainName 
    {
        get
        {
            if (!string.IsNullOrEmpty(UserName) && !string.IsNullOrEmpty(Domain))
            {
                return Domain + "\\" + UserName;
            }
            else
            {
                return "";
            }
        }
    }
}

[DataObject(true)]
public class TerminalSessionsWrapper
{    
    public IEnumerable<TerminalSession> GetSessions()
    {        
        var buffer=IntPtr.Zero;
        var count = 0;
        var size = Marshal.SizeOf(typeof(WTS_SESSION_INFO));        
        TerminalSessionsNative.WTSEnumerateSessions(TerminalSessionsNative.CurrentServer, 0, 1, ref buffer, ref count);        

        var current = (int)buffer;
        for (var i = 0; i < count; i++)
        {
            var si = (WTS_SESSION_INFO)Marshal.PtrToStructure((IntPtr)current, typeof(WTS_SESSION_INFO));
            current += size;
            yield return GetSession(si);            
        }                
    }

    protected TerminalSession GetSession(WTS_SESSION_INFO sessionInfo)
    {
        var session = new TerminalSession();
        session.Id = sessionInfo.SessionID;

        var sessionId = session.Id;

        var buffer = IntPtr.Zero;
        uint returned = 0;

        //winStationName
        TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.WinStationName, out buffer, out returned);
        session.WinStationName = Marshal.PtrToStringAnsi(buffer);
        TerminalSessionsNative.WTSFreeMemory(buffer);

        //userName
        TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.UserName, out buffer, out returned);
        session.UserName = Marshal.PtrToStringAnsi(buffer);
        TerminalSessionsNative.WTSFreeMemory(buffer);

        //domain
        TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.DomainName, out buffer, out returned);
        session.Domain = Marshal.PtrToStringAnsi(buffer);
        TerminalSessionsNative.WTSFreeMemory(buffer);

        TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.ConnectState, out buffer, out returned);
        session.State = (ConnectionState)Marshal.ReadInt32(buffer);
        TerminalSessionsNative.WTSFreeMemory(buffer);

        //clientProtocolType
        TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.ClientProtocolType, out buffer, out returned);
        ushort type = (ushort)Marshal.ReadInt16(buffer);
        if (type == 2)
        {
            if (session.State!=ConnectionState.Listen)
            {
                if (TerminalSessionsNative.WTSQuerySessionInformation(TerminalSessionsNative.CurrentServer, sessionId, InfoClass.ClientAddress, out buffer, out returned))
                {
                    WTS_CLIENT_ADDRESS addr = (WTS_CLIENT_ADDRESS)Marshal.PtrToStructure(buffer, typeof(WTS_CLIENT_ADDRESS));
                    string s = "";

                    switch (addr.AddressFamily)
                    {
                        case 2:
                            s = "IP";
                            break;
                        case 26:
                            s = "IPv6";
                            break;
                        case 6:
                            s = "IPX";
                            break;
                        case 17:
                            s = "NetBIOS";
                            break;
                        case 0:
                            s = "IP";
                            break;
                    }
                    session.Address = string.Format("{4}: {0}.{1}.{2}.{3}", addr.Address[2], addr.Address[3], addr.Address[4], addr.Address[5], s);
                } 
            }
        }
        TerminalSessionsNative.WTSFreeMemory(buffer);       

        return session;
    }
}
